#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Wed Jul  4 23:02:20 2018
@project: 天池比赛-A股主板上市公司公告信息抽取
@group: MZH_314
@author: LHQ
"""
import os
import re
from array import array

import numpy as np
import pandas as pd
from sklearn import svm
from sklearn.model_selection import KFold, cross_val_score
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_selection import SelectKBest
from sklearn.pipeline import Pipeline
from sklearn.feature_selection import chi2
from sklearn.externals import joblib

from reportIE.utils.segment import seg


def process_fields(fields):
    field_words = []
    for field in fields:
        field_words.append(" ".join([w for w in seg(field) if w.strip()]))
    return field_words


class ColFeature:
    def __init__(self, field_feat):
        self.field_feat = field_feat
    
    @staticmethod
    def build_value_feature(values_list):
        features = []
        for values in values_list:
            f = array('i', [0 for _ in range(5)])
            if np.any([bool(re.search("\d+-\d+", w)) for w in values]):
                f[0] = 1
            if np.any([bool(re.search("^[\d\.]+$", w)) for w in values]):
                f[1] = 1
            if np.any([bool(re.search("\w{2,}", w)) for w in values]):
                f[2] = 1
            if np.any([w.startswith("0.") for w in values]):
                f[3] = 1
            lens = [len(w) for w in values]
            lenlt4 = len([l for l in lens if l < 5])
            lengt4 = len([l for l in lens if l > 5])
            f[4] = 1 if lenlt4 > lengt4 else 0
            features.append(f)
        return np.array(features)
        
    def build_feature(self, fields, values, *args):
        field_words = process_fields(fields)
        features_field = self.field_feat.transform(field_words).toarray()
        features_value = self.build_value_feature(values)
        features = np.concatenate((features_field, features_value, *args), axis=1)
        return features
    
    @classmethod
    def from_modelfile(cls, field_feat_path=None):
        if field_feat_path is None:
            field_feat_path = "models/zengjianchi/table_feat.m"
        field_feat = joblib.load(field_feat_path)
        feat = cls(field_feat)
        return feat


if __name__ == "__main__":
    path = os.path.abspath("../data/training_data/zengjianchi_table.csv")
    
    df = pd.read_csv(path)
  
    fields = df['field'].fillna("").map(str).tolist()
    values = df['values'].map(eval).tolist()
    field_index = df['field_index'].values.reshape((-1, 1))

    y = df['label'].tolist()
      
    fields_word = process_fields(fields)
    
    vectorizer = CountVectorizer()
    vectorizer.fit_transform(fields_word)
    
    tmp_feature = vectorizer.transform(fields_word)
    
    select = SelectKBest(chi2, k=150)
    select.fit_transform(tmp_feature, y)     

    pipe = Pipeline([("vectorizer", vectorizer), ("select", select)])
    
    col_feat = ColFeature(pipe)
#    features = col_feat.build_feature(fields, values, field_index)
    features = col_feat.build_feature(fields, values)
    
    # train
    clf = svm.SVC()
    
    k_fold = KFold(n_splits=5)
    scores = cross_val_score(clf, features, y, cv=k_fold, n_jobs=-1)
    print(scores)
    
    clf.fit(features, y)
    
    # save models
    joblib.dump(clf, "models/zengjianchi/table_clf.m")

    joblib.dump(pipe, "models/zengjianchi/table_feat.m")

